#!/usr/bin/python3

import threading
import os, time
import requests, json
import queue
import datetime

# from scripts.post2audio import main as p2a
from scripts.sovits import Enable_Sound_TTS, Make_Sound_TTs
from scripts.qiniu import qiMainThead
import wave
import numpy as np
import logging


def beijing(sec, what):
    beijing_time = datetime.datetime.now() + datetime.timedelta(hours=8)
    return beijing_time.timetuple()


logging.Formatter.converter = beijing

logger = logging.getLogger(__name__)
logger.setLevel(level=logging.INFO)
handler = logging.FileHandler("../logs/GPTChat.log")
handler.setLevel(logging.INFO)
formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
# formatter.converter = beijing
handler.setFormatter(formatter)
logger.addHandler(handler)

# 禁用安全请求警告
requests.packages.urllib3.disable_warnings(
    requests.packages.urllib3.exceptions.InsecureRequestWarning
)


def remove_file(file_path):
    """移除本地文件"""
    if os.path.exists(file_path):
        os.remove(file_path)
        logger.info(f"移除文件: {file_path}")


def hasConf():
    """获取用户配置"""
    try:
        with open("conf.json", "r") as fs:
            return json.load(fs)
    except:
        return {}


CONF = hasConf()
if not CONF:
    logger.error("请先配置conf.json")
    raise Exception("请先配置conf.json")


class hasTasksRequest(requests.Session):
    host = CONF.get("host", "http://172.22.80.139")

    def do_get(self, url="", headers=None, params=None, data=None):
        url = self.host + url
        try:
            r = self.get(url, headers=headers, params=params, data=data, verify=False)
            return r.json() if r and r.status_code == 200 else {}
        except BaseException as err:
            logger.error(f"{err}")
            return {}

    def do_post(self, url="", headers=None, params=None, data=None, json=None):
        url = self.host + url
        try:
            r = self.post(
                url, headers=headers, params=params, data=data, json=json, verify=False
            )
            logger.warning(f"{r.status_code} {r.text}")
            return r.json() # if r and r.status_code == 200 else {}
        except BaseException as err:
            logger.error(f"{err}")
            return {}

    def has_tasksA(self):
        """获取合成任务"""
        url = "/rebot/postinfer/11/"
        headers = {"Content-Type": "application/json"}
        headers["NAME"] = "aniPortrait"
        headers["VERSION"] = "1.0.0"
        res = self.do_get(url, headers=headers)
        return res

    def has_tasksB(self):
        """获取合成任务"""
        url = "/rebot/postinfer/12/"
        headers = {"Content-Type": "application/json"}
        headers["NAME"] = "aniPortrait"
        headers["VERSION"] = "1.0.0"
        res = self.do_get(url, headers=headers)
        return res

    def has_attribute(self, uid):
        """查询AI属性"""
        url = f"/rebot/rebotConf/{uid}/"
        headers = {"Content-Type": "application/json"}
        headers["NAME"] = "aniPortrait"
        headers["VERSION"] = "1.0.0"
        res = self.do_get(url, headers=headers)
        return res

    def has_audio(self, uid):
        """查询音频属性"""
        url = f"/rebot/audioConf/{uid}/"
        headers = {"Content-Type": "application/json"}
        headers["NAME"] = "aniPortrait"
        headers["VERSION"] = "1.0.0"
        res = self.do_get(url, headers=headers)
        return res

    def post_task_refer(self, guid, uid, code, data, mode="Audio", plan="1/1"):
        """提交合成任务"""
        url = "/rebot/postinfer/"
        headers = {"Content-Type": "application/json"}
        headers["NAME"] = "aniPortrait"
        headers["VERSION"] = "1.0.0"
        data = {
            "guid": guid,
            # "user": uid,
            "code": code,
            "data": data,
            "mode": mode,
            "plan": "Audio",
        }
        res = self.do_post(url, headers=headers, data=json.dumps(data))
        logger.info(f"提交任务: {guid} {code} {res}")
        return res


class myThread(threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)
        self.is_Run = True
        self.web = hasTasksRequest()

    def stop(self):
        self.is_Run = False

    def run(self):
        print("myThread Thread is running...")
        logger.debug("myThread Thread is running...")
        while self.is_Run:
            # 音频11
            task = self.web.has_tasksA().get("resmsg")            
            if task:
                # print(type(task), task)              
                # 格式化数据
                t, u, i, d = self.formatMsgA(task)
                print(
                    f'{time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())} 新任务开始:{t} {u} {i} {d}'
                )
                logger.info(f'新的任务:{task["id"]} {t} {u} {i}')  
                # 格式化字符串
                txt = (
                    d.strip()
                    .replace(r"\n\n", "。")
                    .replace(r"\n", "。")
                    .replace("*", "")
                    .replace("#", "")
                    .replace("/", "")
                    .replace(" ", "")
                    .replace("-", "")
                    .replace('"', "")
                    .replace("'", "")
                    .replace("“", "")
                    .replace("”", "")
                    .replace("’", "")
                    .replace("‘", "")
                    .replace(":", "")
                    .replace("：", "")
                    .replace("?", "。")
                    .replace("？", "。")
                    .replace("!", "。")
                    .replace("！", "。")
                    .replace("。。", "。")
                )
                print(
                    time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
                    "格式化文本:\n",
                    txt,
                )

                # 计算文本分块
                idx, chunk, local = [], 24, []
                # 按字数分割
                # for c in range(0, len(txt), chunk):
                #     idx.append(txt[c : c + chunk])
                # 按符号分割
                for c in txt.split("。"):
                    if len(c) > 1:
                        idx.append(f"{c}")
                # 查询音频属性
                au = self.web.has_audio(i).get("resmsg")
                if not au:
                    print("没有找到音频属性")
                    logger.error(f"没有找到音频属性 {au}")
                    self.web.post_task_refer(u, i, 400, "没有找到音频属性", t)
                    continue
                if not au.get("status"):
                    print("音频属性未训练")
                    logger.error(f"音频属性未训练 {au}")
                    self.web.post_task_refer(u, i, 400, "音频属性未训练", t)
                    continue

                # 开启TTS功能
                isEnable = Enable_Sound_TTS(
                    {
                        "gptmodels": au.get("gptmodels"),
                        "sovitsmodels": au.get("sovitsmodels"),
                    },
                    True,
                )
                if not isEnable:
                    print("AI服务器启动失败")
                    logger.error(f"AI服务器启动失败 {isEnable}")
                    self.web.post_task_refer(u, i, 400, "AI服务器启动失败", t)
                    # self.doRestartSoVits()
                    continue

                try:
                    # 开始循环TTS
                    for g in range(len(idx)):
                        print(f"\n{idx[g]}")
                        tts = Make_Sound_TTs(au, idx[g])
                        if not tts:
                            raise AssertionError(f"TTS失败 {idx[g]}")
                        # 加入列表
                        local.append(tts)
                except BaseException as err:
                    print("AI服务器TTS失败")
                    logger.error(f"AI服务器TTS失败 {err}")
                    self.web.post_task_refer(u, i, 400, "AI服务器TTS失败", t)
                    # self.doRestartSoVits()
                    continue

                # 合并文件
                if not local:
                    print("分片编组为空")
                    logger.error(f"分片编组为空 {local}")
                    self.web.post_task_refer(u, i, 400, "分片编组为空", t)
                    continue
                new_arry = np.concatenate([self.hasarrays(n) for n in local])
                self.saveOutAudio(new_arry)
                print("分片合并完成")

                # 上传文件
                r = qiMainThead().doUpload("output.wav")
                print(
                    f'{time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())} 上传七牛云完成 https://oss.3vppt.com/{r.get("key")}'
                )
                self.web.post_task_refer(
                    u,
                    i,
                    200,
                    f'https://oss.3vppt.com/{r["key"]}',
                    t,
                    "1/1",
                    # f"{g+1}/{len(idx)}",
                )
                logger.info(
                    f"任务{u}上传七牛云完成 https://oss.3vppt.com/{r.get('key')}"
                )

                # 关闭TTS功能
                isEnable = Enable_Sound_TTS(
                    {
                        "gptmodels": au.get("gptmodels"),
                        "sovitsmodels": au.get("sovitsmodels"),
                    },
                    False,
                )
                print("任务结束")
                logger.info(f'任务完成: {task["id"]} {u}')
            time.sleep(3)

    def formatMsgB(self, msg):
        """返回: 对象"""
        return eval(msg)

    def formatMsgA(self, msg):
        """返回: 类型, guid, model_id, 内容"""
        r = msg
        return (r.get("mode", "Audio"), r.get("guid"), r.get("model"), r.get("chat"))

    @staticmethod
    def hasarrays(url):
        with wave.open(url, "rb") as f:
            params = f.getparams()
            print(params)
            frames = f.readframes(params.nframes)
            arrays = np.frombuffer(frames, dtype=np.int16)
            return arrays

    @staticmethod
    def saveOutAudio(arrays):
        with wave.open("output.wav", "wb") as f:
            f.setparams((1, 2, 32000, 0, "NONE", "not compressed"))
            f.writeframes(arrays.tobytes())

    @staticmethod
    def downloadAudio(url):
        r = requests.get(url, stream=True)
        with open("audio.wav", "wb") as f:
            for chunk in r.iter_content(chunk_size=1024):
                if chunk:
                    f.write(chunk)
        return "audio.wav"

    @staticmethod
    def downloadImage(url):
        r = requests.get(url, stream=True)
        with open("image.png", "wb") as f:
            for chunk in r.iter_content(chunk_size=1024):
                if chunk:
                    f.write(chunk)
        return "image.png"

    @staticmethod
    def doRestartSoVits():
        """重启AI服务"""
        print("重启AI服务中...")
        with open("aistatus", "w") as f:
            f.write("error")
        for i in range(60):
            print(f"\r{i}", end="")
            time.sleep(1)


# def do_image():
#     """测试视频生成"""
#     audio = "http://oss.3vppt.com/ppt_media/FlS-OsAGXVO9kTV6Fj5xF4qE3eY0&audio.wav"
#     image = "https://oss.3vppt.com/ppt_media/nv01.png"
#     video = p2a(audio, image)
#     print(video)
#     # remove_file(audio["video"])


def do_audio():
    """测试音频生成"""
    audio = "http://oss.3vppt.com/ppt_media/FlS-OsAGXVO9kTV6Fj5xF4qE3eY0&audio.wav"
    text = "你好,有什么可以效劳的吗？"
    # Enable_Sound_TTS, Make_Sound_TTs
    # 开启TTS功能
    isEnable = Enable_Sound_TTS(
        {
            "gptmodels": "mona-e10.ckpt",
            "sovitsmodels": "mona_e15_s1110.pth",
        },
        True,
    )
    # 开始TTS
    content = "你好,今天的天气真好呀，有什么可以效劳的吗？"
    conf = {"sound": audio, "refer": text, "language": "Chinese"}
    tts = Make_Sound_TTs(conf, content)
    print(tts)


if __name__ == "__main__":
    print("Hello World")
    print(
        """
        音频推理  ver:3.2              
        """
    )
    print(f'{"x"*40}')
    [print(k, ": ", v) for k, v in CONF.items()]
    print(f'{"x"*40}\n')
    main = myThread()
    main.start()
    logger.info("main start")
    while True:
        x = input(">")
        if x == "stop":
            main.stop()
            break
        # if x == "start":
        #     main.start()
        # if x == "image":
        #     do_image()
        # if x == "audio":
        #     do_audio()
    main.join()
    logger.info("main stop")
    print("Bye")
